import glob
import os.path

import torch.utils.data as td
import pickle

class MyDataset(td.Dataset):
    def __init__(self, pathFile):

        with open(pathFile, 'rb') as f:
            data = pickle.load(f)
        self.audioInput = data[0]
        self.labelInput = data[1]

    def __len__(self):
        return self.audioInput.shape[0]

    def __getitem__(self, index):
        sinAudio = self.audioInput[index]
        sinLabel = self.labelInput[index]

        sample = {'audioInput':[], 'labelInput':[] ,'audioLen':[], 'labelLen':[]}
        sample['audioInput'] = sinAudio
        sample['labelInput'] = sinLabel
        return sample
    # 最后输出的数据类型是torch

class MultiDataset(td.Dataset):
    def __init__(self, pathFile):
        prefix = pathFile.split('.')[0]
        self.prefix = prefix
        searchPath = os.path.join(prefix, '*.dp')
        self.allPath = []

        for path in glob.glob(searchPath):
            self.allPath.append(path)

    def __len__(self):
        return len(self.allPath)

    def __getitem__(self, index):
        loadPath = os.path.join(self.prefix, f'{index}.dp')
        if loadPath not in self.allPath:
            raise ValueError
        with open(loadPath, 'rb') as f:
            sinAudio, sinLabel = pickle.load(f)

        sample = {'audioInput':[], 'labelInput':[] ,'audioLen':[], 'labelLen':[]}
        sample['audioInput'] = sinAudio
        sample['labelInput'] = sinLabel
        return sample